
import torch
import os, glob
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from torch.utils.data import DataLoader


class classiferDataset(Dataset):
    def __init__(self, imgPath, labelPath):

        self.imgPath = imgPath
        self.labelPath = labelPath

        self.imgsPathList = glob.glob(imgPath + "*.png")
        # self.labels = glob.glob(labelPath + "*.txt")
        # assert self.imgs.__len__() == self.labels.__len__(), print("imgs num not match labels num")
        self.labelsPathList = [labelPath + i.split("\\")[-1][:-3] + "txt" for i in self.imgsPathList]


    def __len__(self):
        return self.imgsPathList.__len__()


    def __getitem__(self, index):
        imgPath, labelPath = self.imgsPathList[index], self.labelsPathList[index]

        tf = transforms.Compose([
            lambda x: Image.open(x).convert('RGB'),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            # transforms.Normalize(mean=[0.485, 0.456, 0.406],
            #                     std=[0.229, 0.224, 0.225])
        ])

        img = tf(imgPath)

        with open(labelPath, "r") as f:
            label = int(f.read()) - 1

        return img, label



class VeRiClassiferDataset(Dataset):

    def __init__(self, imgPath):

        self.imgPath = imgPath

        self.imgsPathList = glob.glob(imgPath + "*.jpg")
        # self.labels = glob.glob(labelPath + "*.txt")
        # assert self.imgs.__len__() == self.labels.__len__(), print("imgs num not match labels num")
        # self.labels = [int(i.split("\\")[-1][:4]) - 1 for i in self.imgsPathList]

        self.labels = []
        self.classNum = 0
        self.idxMap = {}
        for i in self.imgsPathList:
            k = int(i.split("\\")[-1][:4])
            if k in self.idxMap.keys():
                self.labels.append(self.idxMap[k])
            else:
                self.idxMap[k] = self.classNum
                self.labels.append(self.classNum)
                self.classNum += 1



    def __len__(self):
        return self.imgsPathList.__len__()


    def __getitem__(self, index):
        imgPath, label = self.imgsPathList[index], self.labels[index]

        tf = transforms.Compose([
            lambda x: Image.open(x).convert('RGB'),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            # transforms.Normalize(mean=[0.485, 0.456, 0.406],
            #                     std=[0.229, 0.224, 0.225])
        ])

        img = tf(imgPath)

        return img, label



class VeRiCircleDataset(Dataset):

    def __init__(self, imgPath):

        self.imgPath = imgPath

        self.imgsPathList = glob.glob(imgPath + "*.jpg")
        # self.labels = glob.glob(labelPath + "*.txt")
        # assert self.imgs.__len__() == self.labels.__len__(), print("imgs num not match labels num")
        self.labels = [int(i.split("\\")[-1][:4]) - 1 for i in self.imgsPathList]

        # self.labels = []
        # self.classNum = 0
        # self.idxMap = {}
        # for i in self.imgsPathList:
        #     k = int(i.split("\\")[-1][:4])
        #     if k in self.idxMap.keys():
        #         self.labels.append(self.idxMap[k])
        #     else:
        #         self.idxMap[k] = self.classNum
        #         self.labels.append(self.classNum)
        #         self.classNum += 1



    def __len__(self):
        return self.imgsPathList.__len__()


    def __getitem__(self, index):
        imgPath, label = self.imgsPathList[index], self.labels[index]

        tf = transforms.Compose([
            lambda x: Image.open(x).convert('RGB'),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            # transforms.Normalize(mean=[0.485, 0.456, 0.406],
            #                     std=[0.229, 0.224, 0.225])
        ])

        img = tf(imgPath)

        return img, label




if __name__ == '__main__':
    # print(int("0001"))
    ds = VeRiClassiferDataset(".\\data\\VeRi\\image_train\\")
    print(ds.imgsPathList)
    print(ds.labels)
    print(len(ds.imgsPathList))
    print(len(ds.labels))

    # d = { 2 :3 , 4: 5}
    # print(d[2])
    # print(4 in d.keys())
    # print(5 in d.keys())